#ifndef XLCOMM_NET_PROTOBUF_DISPATCHER_H_
#define XLCOMM_NET_PROTOBUF_DISPATCHER_H_

#include <map>
#include <memory>

#include "codec.h"
#include "xlcomm/noncopyable.h"

class Callback : noncopyable {
 public:
  virtual ~Callback() = default;
  virtual void OnMessage(const net::TcpConnectionPtr& conn,
                         const MessagePtr& message,
                         Timestamp recv_time) const = 0;
};

template <typename T>
class TCallback : public Callback {
 public:
  using TMessageCallback = std::function<void(const net::TcpConnectionPtr& conn,
                                              const std::shared_ptr<T>& message,
                                              Timestamp recv_time)>;

  explicit TCallback(const TMessageCallback& cb) : t_msg_cb_(cb) {
    if (std::is_base_of<google::protobuf::Message, T>::value != true)
      LOG_FATAL << "T must be derived from google::protobuf::Message";
  }

  void OnMessage(const net::TcpConnectionPtr& conn, const MessagePtr& message,
                 Timestamp recv_time) const override {
    std::shared_ptr<T> msg = std::dynamic_pointer_cast<T>(message);
    if (t_msg_cb_) t_msg_cb_(conn, msg, recv_time);
  }

 private:
  TMessageCallback t_msg_cb_;
};

void DefaultUnknownMsgCb(const net::TcpConnectionPtr& conn,
                         const MessagePtr& message, Timestamp recv_time) {
  LOG_INFO << "DefaultUnknownMsgCb only shutdown";
  conn->Shutdown();
}

class ProtobufDispatcher {
 public:
  explicit ProtobufDispatcher(
      const ProtobufMessageCallback& cb = DefaultUnknownMsgCb)
      : unknown_msg_cb_(cb) {}

  template <typename T>
  void RegisterMessageCallback(
      const typename TCallback<T>::TMessageCallback& tcb) {
    std::shared_ptr<TCallback<T>> cbptr(new TCallback<T>(tcb));
    messages_[T::descriptor()] = cbptr;
  }

  void OnProtobufMessage(const net::TcpConnectionPtr& conn,
                         const MessagePtr& message, Timestamp recv_time) {
    MessageMap::const_iterator it = messages_.find(message->GetDescriptor());
    if (it != messages_.end()) {
      it->second->OnMessage(conn, message, recv_time);
    } else {
      LOG_WARN << "Unknown message type:" << message->GetTypeName();
      if (unknown_msg_cb_) unknown_msg_cb_(conn, message, recv_time);
    }
  }

 private:
  using MessageMap =
      std::map<const google::protobuf::Descriptor*, std::shared_ptr<Callback>>;

  MessageMap messages_;
  ProtobufMessageCallback unknown_msg_cb_;
};

#endif  // XLCOMM_NET_PROTOBUF_DISPATCHER_H_